#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
智谱文本大模型分析器模块
负责调用智谱文本大模型API进行补充分析
"""

import os
import sys
import json
import logging
import time
import requests
from typing import Dict, Any, List, Optional, Union
from datetime import datetime

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)


class ZhipuTextAnalyzer:
    """
    智谱文本大模型分析器
    封装智谱AI文本模型API调用
    """
    
    # 默认配置
    DEFAULT_API_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    DEFAULT_MODEL = "glm-4"
    DEFAULT_TIMEOUT = 60
    MAX_RETRIES = 3
    
    def __init__(self, 
                 api_key: str = None,
                 api_url: str = None,
                 model: str = DEFAULT_MODEL,
                 timeout: int = DEFAULT_TIMEOUT,
                 config_path: str = None):
        """
        初始化智谱文本分析器
        
        Args:
            api_key: 智谱API密钥
            api_url: API调用地址
            model: 使用的模型名称
            timeout: 请求超时时间
            config_path: 配置文件路径
        """
        # 加载配置
        self._load_config(config_path)
        
        # 设置API密钥
        self._api_key = api_key or self._config.get("api_key")
        if not self._api_key:
            # 尝试从环境变量获取
            self._api_key = os.environ.get("ZHIPU_API_KEY")
            
            # 尝试从.env文件获取
            if not self._api_key:
                env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), '.env')
                if os.path.exists(env_path):
                    try:
                        with open(env_path, 'r', encoding='utf-8') as f:
                            for line in f:
                                if line.startswith('ZHIPU_API_KEY='):
                                    self._api_key = line.split('=', 1)[1].strip().strip('"')
                                    break
                    except Exception as e:
                        logger.warning(f"读取.env文件失败: {str(e)}")
        
        # 设置其他参数
        self._api_url = api_url or self._config.get("api_url", self.DEFAULT_API_URL)
        self._model = model or self._config.get("model", self.DEFAULT_MODEL)
        self._timeout = timeout
        self._session = requests.Session()
        
        logger.info(f"智谱文本分析器初始化完成，模型: {self._model}")
    
    def _load_config(self, config_path: str = None):
        """
        加载配置文件
        
        Args:
            config_path: 配置文件路径
        """
        self._config = {}
        
        if config_path and os.path.exists(config_path):
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    self._config = json.load(f)
                logger.info(f"配置文件加载成功: {config_path}")
            except Exception as e:
                logger.error(f"加载配置文件失败: {str(e)}")
        else:
            # 尝试加载默认配置
            default_config = os.path.join(os.path.dirname(__file__), "config.json")
            if os.path.exists(default_config):
                try:
                    with open(default_config, 'r', encoding='utf-8') as f:
                        self._config = json.load(f)
                except Exception as e:
                    logger.warning(f"加载默认配置失败: {str(e)}")
    
    def analyze_text(self, 
                    prompt: str,
                    system_prompt: str = None,
                    temperature: float = 0.7,
                    max_tokens: int = 2048,
                    retry_on_error: bool = True) -> Dict[str, Any]:
        """
        分析文本内容
        
        Args:
            prompt: 用户提示词
            system_prompt: 系统提示词
            temperature: 生成温度
            max_tokens: 最大生成token数
            retry_on_error: 错误时是否重试
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        # 构建请求消息
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        
        messages.append({"role": "user", "content": prompt})
        
        # 构建请求参数
        payload = {
            "model": self._model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": False
        }
        
        # 发送请求
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self._api_key}"
        }
        
        # 重试机制
        retries = 0
        while retries <= (self.MAX_RETRIES if retry_on_error else 0):
            try:
                logger.info(f"正在调用智谱文本API，重试次数: {retries}")
                response = self._session.post(
                    self._api_url,
                    headers=headers,
                    json=payload,
                    timeout=self._timeout
                )
                
                # 检查响应状态
                response.raise_for_status()
                
                # 解析响应
                result = response.json()
                
                # 检查是否有错误
                if "error" in result:
                    raise Exception(f"API错误: {result['error'].get('message', str(result['error']))}")
                
                # 提取生成的内容
                if "choices" in result and result["choices"]:
                    generated_content = result["choices"][0].get("message", {}).get("content", "")
                    
                    return {
                        "success": True,
                        "content": generated_content,
                        "token_usage": result.get("usage", {}),
                        "model": result.get("model"),
                        "timestamp": datetime.now().isoformat()
                    }
                else:
                    raise Exception("API返回格式异常，没有choices字段")
                    
            except requests.exceptions.RequestException as e:
                logger.error(f"请求异常: {str(e)}")
                if retries >= (self.MAX_RETRIES if retry_on_error else 0):
                    return {"success": False, "error": f"请求失败: {str(e)}"}
                
            except Exception as e:
                logger.error(f"分析失败: {str(e)}")
                if retries >= (self.MAX_RETRIES if retry_on_error else 0):
                    return {"success": False, "error": str(e)}
            
            retries += 1
            if retries <= (self.MAX_RETRIES if retry_on_error else 0):
                wait_time = min(2 ** retries, 30)  # 指数退避，最大等待30秒
                logger.info(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
    
    def analyze_video_encoding(self, 
                             visual_analysis: Dict[str, Any],
                             video_info: Dict[str, Any] = None,
                             low_confidence_threshold: float = 60.0,
                             categories_to_analyze: List[str] = None) -> Dict[str, Any]:
        """
        分析视频编码类目，补充视觉分析中置信度低的字段
        
        Args:
            visual_analysis: 视觉分析结果
            video_info: 视频信息
            low_confidence_threshold: 低置信度阈值
            categories_to_analyze: 要分析的类目列表，如果为None则分析所有低置信度类目
            
        Returns:
            Dict[str, Any]: 补充分析结果
        """
        # 准备分析数据
        prompt, system_prompt = self._prepare_encoding_prompt(
            visual_analysis,
            video_info,
            low_confidence_threshold,
            categories_to_analyze
        )
        
        if not prompt:
            logger.info("没有需要补充分析的类目")
            return {
                "success": True,
                "message": "没有需要补充分析的类目",
                "analyzed_categories": []
            }
        
        # 调用文本分析
        result = self.analyze_text(prompt, system_prompt)
        
        if not result["success"]:
            return result
        
        # 解析结果
        try:
            parsed_result = self._parse_encoding_result(result["content"])
            
            # 合并结果
            final_result = {
                "success": True,
                "analyzed_categories": parsed_result["categories"],
                "full_response": result["content"],
                "token_usage": result.get("token_usage"),
                "analysis_time": datetime.now().isoformat()
            }
            
            # 如果有需要补充的类目，添加到结果
            if parsed_result["categories"]:
                final_result["recommended_updates"] = self._generate_recommended_updates(
                    visual_analysis,
                    parsed_result["categories"]
                )
            
            return final_result
            
        except Exception as e:
            logger.error(f"解析结果失败: {str(e)}")
            return {
                "success": False,
                "error": f"解析结果失败: {str(e)}",
                "raw_response": result["content"]
            }
    
    def _prepare_encoding_prompt(self, 
                               visual_analysis: Dict[str, Any],
                               video_info: Dict[str, Any],
                               low_confidence_threshold: float,
                               categories_to_analyze: List[str]) -> Tuple[str, str]:
        """
        准备编码分析提示词
        
        Args:
            visual_analysis: 视觉分析结果
            video_info: 视频信息
            low_confidence_threshold: 低置信度阈值
            categories_to_analyze: 要分析的类目列表
            
        Returns:
            Tuple[str, str]: (用户提示词, 系统提示词)
        """
        # 提取需要分析的类目
        categories_to_process = []
        encoding_categories = visual_analysis.get("encoding_categories", {})
        
        for category_name, category_data in encoding_categories.items():
            # 检查是否在指定的分析列表中
            if categories_to_analyze and category_name not in categories_to_analyze:
                continue
            
            # 检查置信度
            confidence = category_data.get("confidence", 0)
            if confidence < low_confidence_threshold:
                categories_to_process.append({
                    "name": category_name,
                    "current_option": category_data.get("option", "未知"),
                    "current_confidence": confidence,
                    "current_reasoning": category_data.get("reasoning", "无")
                })
        
        # 如果没有需要处理的类目，返回空
        if not categories_to_process:
            return "", ""
        
        # 构建系统提示词
        system_prompt = """
你是一位专业的视频内容分析专家，负责根据视觉分析结果和视频信息，对抖音视频进行编码类目补充分析。
请严格按照输出格式要求返回JSON格式的分析结果。

分析原则：
1. 基于提供的视觉分析结果进行推理补充
2. 对于置信度低的字段，提供更准确的判断和更高的置信度
3. 对于需要文本理解的类目（如主题类型、情感倾向等），提供深入分析
4. 每个分析必须有明确的推理过程

输出格式：
```json
{
  "categories": [
    {
      "name": "类目名称",
      "option": "最佳匹配选项",
      "confidence": 置信度（0-100）,
      "reasoning": "详细的推理过程"
    }
  ]
}
```
请确保返回的JSON格式正确，不要包含任何JSON之外的文本。
        """.strip()
        
        # 构建用户提示词
        prompt_parts = []
        prompt_parts.append("请帮我对以下抖音视频的编码类目进行补充分析，重点关注置信度较低的字段：\n")
        
        # 添加视频信息
        if video_info:
            prompt_parts.append("【视频基本信息】")
            prompt_parts.append(f"- 时长: {video_info.get('duration', 0):.2f}秒")
            prompt_parts.append(f"- 分辨率: {video_info.get('width', 0)}x{video_info.get('height', 0)}")
            prompt_parts.append(f"- 帧率: {video_info.get('fps', 0):.2f}fps")
            prompt_parts.append("")
        
        # 添加需要分析的类目
        prompt_parts.append("【需要补充分析的类目】")
        for category in categories_to_process:
            prompt_parts.append(f"类目: {category['name']}")
            prompt_parts.append(f"- 当前选项: {category['current_option']}")
            prompt_parts.append(f"- 当前置信度: {category['current_confidence']}%")
            prompt_parts.append(f"- 当前推理: {category['current_reasoning']}")
            prompt_parts.append("")
        
        # 添加视觉分析摘要
        visual_summary = self._generate_visual_summary(visual_analysis)
        if visual_summary:
            prompt_parts.append("【视觉分析摘要】")
            prompt_parts.append(visual_summary)
            prompt_parts.append("")
        
        # 添加指令
        prompt_parts.append("请根据上述信息，对置信度低的类目进行补充分析，并按照指定格式返回JSON结果。")
        
        prompt = "\n".join(prompt_parts)
        
        return prompt, system_prompt
    
    def _generate_visual_summary(self, visual_analysis: Dict[str, Any]) -> str:
        """
        生成视觉分析摘要
        
        Args:
            visual_analysis: 视觉分析结果
            
        Returns:
            str: 视觉分析摘要
        """
        parts = []
        
        # 提取关键信息
        key_frames = visual_analysis.get("keyframe_count", 0)
        overall_confidence = visual_analysis.get("overall_confidence", 0)
        
        parts.append(f"- 分析关键帧数量: {key_frames}")
        parts.append(f"- 总体置信度: {overall_confidence:.1f}%")
        
        # 提取高置信度的分析结果
        encoding_categories = visual_analysis.get("encoding_categories", {})
        high_confidence_categories = []
        
        for category_name, category_data in encoding_categories.items():
            if category_data.get("confidence", 0) >= 80:
                high_confidence_categories.append({
                    "name": category_name,
                    "option": category_data.get("option", "未知"),
                    "confidence": category_data.get("confidence", 0)
                })
        
        if high_confidence_categories:
            parts.append("\n高置信度分析结果:")
            for category in high_confidence_categories:
                parts.append(f"  - {category['name']}: {category['option']} ({category['confidence']}%)")
        
        # 提取补充信息
        if "complementary_info" in visual_analysis:
            complementary = visual_analysis["complementary_info"]
            if complementary:
                parts.append("\n视频补充信息:")
                for key, value in complementary.items():
                    if isinstance(value, str) and len(value) > 50:
                        value = value[:50] + "..."
                    parts.append(f"  - {key}: {value}")
        
        return "\n".join(parts)
    
    def _parse_encoding_result(self, response_content: str) -> Dict[str, Any]:
        """
        解析编码分析结果
        
        Args:
            response_content: API返回的原始内容
            
        Returns:
            Dict[str, Any]: 解析后的结果
        """
        # 尝试提取JSON部分
        import re
        
        # 查找JSON格式的内容
        json_match = re.search(r'```json\s*([\s\S]*?)\s*```', response_content)
        
        if json_match:
            json_str = json_match.group(1)
        else:
            # 尝试直接解析整个内容
            json_str = response_content
        
        # 解析JSON
        try:
            result = json.loads(json_str)
            if "categories" not in result:
                result["categories"] = []
            return result
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败: {str(e)}")
            # 尝试使用备用解析方法
            return self._parse_fallback(response_content)
    
    def _parse_fallback(self, response_content: str) -> Dict[str, Any]:
        """
        备用解析方法
        
        Args:
            response_content: API返回的原始内容
            
        Returns:
            Dict[str, Any]: 解析后的结果
        """
        categories = []
        
        # 简单的正则匹配
        import re
        
        # 尝试匹配类目格式
        category_patterns = [
            r'类目[:：]\s*(\S+)\s+选项[:：]\s*(\S+)\s+置信度[:：]\s*(\d+)%',
            r'"name"\s*[:：]\s*"([^"]+)".*?"option"\s*[:：]\s*"([^"]+)".*?"confidence"\s*[:：]\s*(\d+)',
            r'(\S+)[:：]\s*(\S+)\s+\((\d+)%\)'  # 匹配如 "场景设定: 校园场景 (95%)"
        ]
        
        for pattern in category_patterns:
            matches = re.finditer(pattern, response_content, re.MULTILINE | re.DOTALL)
            for match in matches:
                if len(match.groups()) >= 3:
                    category = {
                        "name": match.group(1).strip(),
                        "option": match.group(2).strip(),
                        "confidence": int(match.group(3).strip()),
                        "reasoning": "基于备用解析方法"
                    }
                    # 检查是否已存在
                    if not any(c["name"] == category["name"] for c in categories):
                        categories.append(category)
        
        return {"categories": categories}
    
    def _generate_recommended_updates(self, 
                                    visual_analysis: Dict[str, Any],
                                    analyzed_categories: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        生成推荐的更新
        
        Args:
            visual_analysis: 原始视觉分析结果
            analyzed_categories: 分析后的类目列表
            
        Returns:
            Dict[str, Any]: 推荐的更新
        """
        updates = {}
        original_categories = visual_analysis.get("encoding_categories", {})
        
        for analyzed in analyzed_categories:
            category_name = analyzed["name"]
            
            # 检查是否有改进
            if category_name in original_categories:
                original_confidence = original_categories[category_name].get("confidence", 0)
                new_confidence = analyzed["confidence"]
                
                # 如果置信度提高了，建议更新
                if new_confidence > original_confidence + 10:  # 置信度提升至少10%才建议更新
                    updates[category_name] = {
                        "old_option": original_categories[category_name].get("option", "未知"),
                        "old_confidence": original_confidence,
                        "new_option": analyzed["option"],
                        "new_confidence": new_confidence,
                        "improvement": new_confidence - original_confidence,
                        "reasoning": analyzed["reasoning"]
                    }
            else:
                # 新增类目
                updates[category_name] = {
                    "action": "add",
                    "option": analyzed["option"],
                    "confidence": analyzed["confidence"],
                    "reasoning": analyzed["reasoning"]
                }
        
        return updates
    
    def analyze_specific_category(self, 
                                category_name: str,
                                category_options: List[str],
                                context_description: str,
                                visual_cues: str = "") -> Dict[str, Any]:
        """
        分析特定的编码类目
        
        Args:
            category_name: 类目名称
            category_options: 可选选项列表
            context_description: 上下文描述
            visual_cues: 视觉提示信息
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        # 构建系统提示词
        system_prompt = f"""
你是一位专业的视频内容分类专家，负责根据提供的信息确定视频的"{category_name}"类目。
请从给定的选项中选择最合适的一个，并提供详细的推理过程。

分析要求：
1. 严格从提供的选项列表中选择
2. 给出明确的置信度评分（0-100）
3. 提供详细的推理理由

输出格式：
```json
{
  "option": "选择的选项",
  "confidence": 置信度评分,
  "reasoning": "详细的推理过程"
}
```
请确保返回的JSON格式正确，不要包含任何JSON之外的文本。
        """.strip()
        
        # 构建用户提示词
        prompt_parts = []
        prompt_parts.append(f"请帮我确定以下视频的'{category_name}'类目：\n")
        
        # 添加上下文信息
        if context_description:
            prompt_parts.append("【上下文描述】")
            prompt_parts.append(context_description)
            prompt_parts.append("")
        
        # 添加视觉提示
        if visual_cues:
            prompt_parts.append("【视觉提示】")
            prompt_parts.append(visual_cues)
            prompt_parts.append("")
        
        # 添加选项列表
        prompt_parts.append("【可选选项】")
        for i, option in enumerate(category_options, 1):
            prompt_parts.append(f"{i}. {option}")
        
        prompt_parts.append("")
        prompt_parts.append(f"请从上述选项中为该视频选择最合适的'{category_name}'类目。")
        
        prompt = "\n".join(prompt_parts)
        
        # 调用文本分析
        result = self.analyze_text(prompt, system_prompt)
        
        if not result["success"]:
            return result
        
        # 解析结果
        try:
            parsed_result = self._parse_category_result(result["content"])
            
            # 添加类目名称
            final_result = {
                "category_name": category_name,
                "option": parsed_result["option"],
                "confidence": parsed_result["confidence"],
                "reasoning": parsed_result["reasoning"],
                "analysis_time": datetime.now().isoformat()
            }
            
            return final_result
            
        except Exception as e:
            logger.error(f"解析特定类目结果失败: {str(e)}")
            return {
                "success": False,
                "error": f"解析结果失败: {str(e)}",
                "raw_response": result["content"]
            }
    
    def _parse_category_result(self, response_content: str) -> Dict[str, Any]:
        """
        解析特定类目分析结果
        
        Args:
            response_content: API返回的原始内容
            
        Returns:
            Dict[str, Any]: 解析后的结果
        """
        # 尝试提取JSON部分
        import re
        
        # 查找JSON格式的内容
        json_match = re.search(r'```json\s*([\s\S]*?)\s*```', response_content)
        
        if json_match:
            json_str = json_match.group(1)
        else:
            # 尝试直接解析整个内容
            json_str = response_content
        
        # 解析JSON
        try:
            result = json.loads(json_str)
            # 验证必要字段
            required_fields = ["option", "confidence", "reasoning"]
            for field in required_fields:
                if field not in result:
                    raise ValueError(f"缺少必要字段: {field}")
            return result
        except (json.JSONDecodeError, ValueError) as e:
            logger.error(f"JSON解析失败: {str(e)}")
            # 尝试使用备用解析方法
            return self._parse_category_fallback(response_content)
    
    def _parse_category_fallback(self, response_content: str) -> Dict[str, Any]:
        """
        特定类目备用解析方法
        
        Args:
            response_content: API返回的原始内容
            
        Returns:
            Dict[str, Any]: 解析后的结果
        """
        import re
        
        # 尝试提取信息
        option_match = re.search(r'选项[:：]\s*(\S+)|"option"\s*[:：]\s*"([^"]+)"', response_content)
        confidence_match = re.search(r'置信度[:：]\s*(\d+)%?|"confidence"\s*[:：]\s*(\d+)', response_content)
        reasoning_match = re.search(r'推理[:：]\s*(.+?)(?:\n|$)', response_content, re.DOTALL)
        
        # 提取选项
        option = "未知"
        if option_match:
            option = option_match.group(1) or option_match.group(2) or "未知"
        
        # 提取置信度
        confidence = 50
        if confidence_match:
            try:
                confidence = int(confidence_match.group(1) or confidence_match.group(2))
                # 确保在有效范围内
                confidence = max(0, min(100, confidence))
            except:
                pass
        
        # 提取推理
        reasoning = "基于备用解析方法"
        if reasoning_match:
            reasoning = reasoning_match.group(1).strip()
        
        return {
            "option": option.strip(),
            "confidence": confidence,
            "reasoning": reasoning
        }
    
    def close(self):
        """
        关闭会话
        """
        if hasattr(self, '_session'):
            self._session.close()
            logger.info("智谱文本分析器会话已关闭")


# 使用示例
if __name__ == '__main__':
    try:
        print("智谱文本分析器初始化中...")
        
        # 创建实例
        analyzer = ZhipuTextAnalyzer()
        
        print("\n智谱文本分析器初始化成功!")
        print("\n可用功能:")
        print("- analyze_text(): 基础文本分析")
        print("- analyze_video_encoding(): 视频编码类目补充分析")
        print("- analyze_specific_category(): 特定类目分析")
        
        # 测试API密钥状态
        has_api_key = analyzer._api_key is not None
        print(f"\nAPI密钥加载状态: {'已加载' if has_api_key else '未加载'}")
        
        if not has_api_key:
            print("\n警告: 请设置ZHIPU_API_KEY环境变量或在.env文件中配置API密钥")
        
    except Exception as e:
        print(f"初始化失败: {str(e)}")
        sys.exit(1)